import csv
import matplotlib.pyplot as plt
import numpy as np
import os
from mpl_toolkits.mplot3d import Axes3D

def plot_3d_flight_data():
    """
    绘制3D飞行轨迹和参考路径的函数
    """
    # 文件路径
    flight_data_file = "results//flight_data.csv"
    waypoints_file = "results//waypoints.csv"
    # reference_path_file = "results//reference_path.csv"

    # 检查文件是否存在
    if not os.path.exists(flight_data_file):
        print(f"Error: File '{flight_data_file}' not found.")
        return
    if not os.path.exists(waypoints_file):
        print(f"Error: File '{waypoints_file}' not found.")
        return
    # if not os.path.exists(reference_path_file):
    #     print(f"Error: File '{reference_path_file}' not found.")
    #     return

    # 读取飞行数据
    flight_data = []
    with open(flight_data_file, 'r') as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            flight_data.append({
                "initial_position": (
                    float(row['Initial Long']),
                    float(row['Initial Lat']),
                    float(row['Initial Alt'])
                ),
                "target_position": (
                    float(row['Target Long']),
                    float(row['Target Lat']),
                    float(row['Target Alt'])
                ),
                "current_position": (
                    float(row['Current Long']),
                    float(row['Current Lat']),
                    float(row['Current Alt'])
                )
            })

    if not flight_data:
        print("No flight data to plot.")
        return

    # 提取飞行轨迹点
    trajectory_longs, trajectory_lats, trajectory_alts = zip(*[data["current_position"] for data in flight_data])

    # 提取初始位置和目标位置
    initial_position = flight_data[0]["initial_position"]
    target_position = flight_data[0]["target_position"]

    # 读取中间检查点
    waypoints = []
    with open(waypoints_file, 'r') as csvfile:
        reader = csv.DictReader(csvfile)
        for row in reader:
            waypoints.append((
                float(row['Long']),
                float(row['Lat']),
                float(row['Alt'])
            ))

    # 提取检查点
    waypoint_longs, waypoint_lats, waypoint_alts = zip(*waypoints) if waypoints else ([], [], [])

    # # 读取参考路径
    # reference_path = []
    # with open(reference_path_file, 'r') as csvfile:
    #     reader = csv.DictReader(csvfile)
    #     for row in reader:
    #         reference_path.append((
    #             float(row['Long']),
    #             float(row['Lat']),
    #             float(row['Alt'])  # 增加高度信息
    #         ))

    # 提取参考路径点
    # ref_longs, ref_lats, ref_alts = zip(*reference_path) if reference_path else ([], [], [])

    # 创建3D图
    fig = plt.figure(figsize=(12, 8))
    ax = fig.add_subplot(111, projection='3d')

    # 绘制参考路径（蓝色线）
    # if ref_longs and ref_lats and ref_alts:
    #     ax.plot(ref_longs, ref_lats, ref_alts, 'b-', label='Reference Path')

    # 绘制飞行轨迹（黑色线）
    ax.plot(trajectory_longs, trajectory_lats, trajectory_alts, 'k-', label='Flight Trajectory')

    # 绘制起点（红色点）
    ax.scatter(
        [initial_position[0]], [initial_position[1]], [initial_position[2]],
        color='red', s=100, label='Start Point'
    )

    # 绘制目标点（绿色点）
    ax.scatter(
        [target_position[0]], [target_position[1]], [target_position[2]],
        color='green', s=100, label='Target Point'
    )

    # 绘制中间检查点（黄色点）
    if waypoint_longs and waypoint_lats and waypoint_alts:
        ax.scatter(
            waypoint_longs, waypoint_lats, waypoint_alts,
            color='yellow', s=100, label='Waypoints'
        )

    # 设置坐标轴标签和标题
    ax.set_xlabel('Longitude')
    ax.set_ylabel('Latitude')
    ax.set_zlabel('Altitude (m)')
    ax.set_title('3D Flight Trajectory and Reference Path')
    ax.legend()
    ax.grid(True)

    # 显示图形
    plt.tight_layout()
    plt.show()

    # 保存图像
    plt.savefig('results//3d_flight_trajectory.png')
    plt.close()
    print("3D Flight trajectory saved as '3d_flight_trajectory.png'")

if __name__ == "__main__":
    plot_3d_flight_data()